

#pragma once

#ifndef YOLODNN
#define YOLODNN

#include <iostream>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;
using namespace cv::dnn;
#include "utils.h"

class datap
{
public:
    datap() {}
    std::vector<int> classIds;      //结果id数组
    std::vector<float> confidences; //结果每个id对应置信度数组
    std::vector<cv::Rect> boxes;    //每个id矩形框
    vector<int> nms_result;
    void reset();
    void press(float *pdata, size_t dlen = 61200, int net_width = 85, float conf_thres = 0.25, float nmsconf = 0.25, float *sc = 0);
    void print();
    void getbox(vector<Output> &output);
    void getbox(vector<Outputf> &output);
};

class yolov5
{
public:
    bool yolov5_P6;
    cv::dnn::Net net;
    yolov5(string netPath = "aa", float _thresh = 0.5);
    ~yolov5();
    bool operator()(cv::Mat &SrcImg, std::vector<Output> &output, float _thresh);
    void drawPred(cv::Mat &img, std::vector<Output> result, std::vector<cv::Scalar> color);
    bool DEBUG;
    datap *dp;
    int cuda;

private:
    int netWidth = 640;  // ONNX图片输入宽度
    int netHeight = 640; // ONNX图片输入高度

    const float netStride[4] = {8, 16.0, 32, 64};

    float boxThreshold = 0.25;
    float classThreshold = 0.2;
    float thresh;
    float nmsThreshold = 0.45;
    float nmsScoreThreshold = boxThreshold * classThreshold;
    int class_num = 80;

public:
    std::vector<std::string> className = {"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
                                          "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
                                          "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
                                          "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
                                          "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
                                          "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
                                          "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
                                          "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
                                          "hair drier", "toothbrush"};
};
#endif